﻿#region COPYRIGHT© 2009-2012 Phillip Clark. All rights reserved.

// For licensing information see License.txt (MIT style licensing).

#endregion

using System;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Diagnostics.Contracts;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using FlitBit.Core;
using FlitBit.Core.Parallel;

namespace FlitBit.Proto
{
	/// <summary>
	///   UDP endpoint for communicating according to a network protocol.
	/// </summary>
	/// <typeparam name="TMessage">message type M</typeparam>
	public class UdpProtocolEndpoint<TMessage> : Disposable
		where TMessage : class
	{
		readonly ConcurrentQueue<Action<byte[]>> _bufferWaiters = new ConcurrentQueue<Action<byte[]>>();
		readonly INetworkProtocol<TMessage> _protocol;
		EndPoint _anyEndPoint;
		Reactor<byte[]> _buffers;
		Reactor<InboundMessage> _notifier;
		Reactor<OverlappedOpState> _receiver;
		INetEndpoint _socket;

		/// <summary>
		///   Creates a new instance.
		/// </summary>
		/// <param name="protocol">a protocol</param>
		public UdpProtocolEndpoint(INetworkProtocol<TMessage> protocol)
		{
			_protocol = protocol;
		}

		/// <summary>
		///   Indicates whether the endpoint is actively receiving.
		/// </summary>
		public bool IsReceiving { get { return _receiver != null; } }

		/// <summary>
		///   Gets the endpoint's protocol.
		/// </summary>
		public INetworkProtocol<TMessage> Protocol { get { return _protocol; } }

		/// <summary>
		///   Begins receiving protocol messages in parallel from the endpoint given.
		/// </summary>
		/// <param name="endpoint">an endpoint</param>
		/// <returns>this object (for chaining)</returns>
		public UdpProtocolEndpoint<TMessage> ParallelReceive(IPEndPoint endpoint)
		{
			Contract.Requires<InvalidOperationException>(!IsReceiving);

			_anyEndPoint = (endpoint.AddressFamily == AddressFamily.InterNetworkV6)
				? new IPEndPoint(IPAddress.IPv6Any, 0)
				: new IPEndPoint(IPAddress.Any, 0);

			_receiver = new Reactor<OverlappedOpState>(Background_Recieve);
			_socket = new OverlappedSocketEndpoint();
			Contract.Assume(!_socket.IsInitialized);
			Contract.Assume(!_socket.IsDisposed);
			_socket.Initialize(CreateSocketWithBind(endpoint));
			Contract.Assume(_socket != null);
			_protocol.Initialize(_socket);
			_notifier = new Reactor<InboundMessage>(Background_Notify);
			_buffers = new Reactor<byte[]>(Background_BufferAvailable);
			Background_ParallelReceiveFromAsync(null);
			return this;
		}

		/// <summary>
		/// Trace message sink. Invokes Core's LogSink.
		/// </summary>
		/// <param name="evType"></param>
		/// <param name="message"></param>
		protected virtual void OnTraceMessage(TraceEventType evType, string message)
		{
			if (LogSink.ShouldTrace(this, evType))
			{
				LogSink.OnTraceEvent(this, evType, message);
			}
		}

		/// <summary>
		///   Recycles the buffer given.
		/// </summary>
		/// <param name="buffer">the buffer</param>
		protected virtual void RecycleBuffer(byte[] buffer)
		{
			_buffers.Push(buffer);
		}

		/// <summary>
		///   Releases a buffer previously acquired from the endpoint.
		/// </summary>
		/// <param name="buffer">the buffer being released</param>
		protected virtual void ReleaseBuffer(byte[] buffer)
		{}

		/// <summary>
		///   Tries to acquire a buffer.
		/// </summary>
		/// <param name="buffer">variable where the buffer will be returned upon success</param>
		/// <returns>true if successful; othewise false</returns>
		protected virtual bool TryAcquireBuffer(out byte[] buffer)
		{
			// default implementation just allocates the buffer
			buffer = new byte[Constants.DefaultReceiveBufferLength];
			return true;
		}

		/// <summary>
		///   Performs the instance's disposal.
		/// </summary>
		/// <param name="disposing"></param>
		/// <returns></returns>
		protected override bool PerformDispose(bool disposing)
		{
			_socket.Dispose();
			_socket = null;

			return disposing;
		}

		/// <summary>
		/// Creates a socket on the endpoint given.
		/// </summary>
		/// <param name="endpoint"></param>
		/// <returns></returns>
		protected virtual Socket CreateSocket(IPEndPoint endpoint)
		{
			var sock = new Socket(endpoint.AddressFamily, SocketType.Dgram, ProtocolType.IP);
			sock.SetSocketOption(SocketOptionLevel.Udp, SocketOptionName.NoDelay, 1);
			sock.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, 1);
			return sock;
		}

		/// <summary>
		/// Creates and binds a socket on the endpoint given.
		/// </summary>
		/// <param name="endpoint"></param>
		/// <returns></returns>
		protected virtual Socket CreateSocketWithBind(IPEndPoint endpoint)
		{
			var sock = new Socket(endpoint.AddressFamily, SocketType.Dgram, ProtocolType.IP);
			sock.SetSocketOption(SocketOptionLevel.Udp, SocketOptionName.NoDelay, 1);
			sock.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, 1);
			sock.Bind(endpoint);
			return sock;
		}

		void Background_BufferAvailable(Reactor<byte[]> reactor, byte[] buffer)
		{
			Action<byte[]> waiter;
			if (_bufferWaiters.TryDequeue(out waiter))
			{
				waiter(buffer);
			}
			else
			{
				ReleaseBuffer(buffer);
			}
		}

		void Background_Notify(Reactor<InboundMessage> reactor, InboundMessage incomming)
		{
			Contract.Requires<ArgumentException>(incomming.Endpoint != null);
			Contract.Requires<ArgumentException>(incomming.Message != null);
			_protocol.HandleEndpointMessage(incomming.Endpoint, incomming.Timestamp, incomming.Message);
		}

		void Background_ParallelReceiveCallback(Exception ex, OverlappedOpState state)
		{
			if (ex == null)
			{
				if (state.SocketError != SocketError.OperationAborted)
				{
					ThreadPool.QueueUserWorkItem(Background_ParallelReceiveFromAsync);
				}
				if (state.SocketError == SocketError.Success)
				{
					if (state.BytesTransferred > 0)
					{
						_receiver.Push(state);
					}
				}
			}
			else
			{
				OnTraceMessage(TraceEventType.Error, String.Concat("Exception during overlapped receive: ", ex.FormatForLogging()));
			}
		}

		void Background_ParallelReceiveFromAsync(object unused)
		{
			if (_socket != null && !_socket.IsDisposed)
			{
				byte[] buffer;
				if (TryAcquireBuffer(out buffer))
				{
					Contract.Assume(!_socket.IsDisposed);
					Contract.Assume(_anyEndPoint != null);
					// Use overlapped IO - the callback will fire when bytes arrive.
					_socket.ReceiveFromAsync(_anyEndPoint, buffer, 0, buffer.Length, Background_ParallelReceiveCallback);
				}
				else
				{
					_bufferWaiters.Enqueue(
																 b =>
																	this._socket.ReceiveFromAsync(this._anyEndPoint, b, 0, b.Length, Background_ParallelReceiveCallback));
				}
			}
		}

		void Background_Recieve(Reactor<OverlappedOpState> reactor, OverlappedOpState op)
		{
			Contract.Requires<ArgumentException>(op.RemoteEndPoint != null);
			Contract.Requires<ArgumentException>(op.Buffer != null);
			Contract.Requires<ArgumentException>(op.BytesTransferred >= 0);
			Contract.Requires<ArgumentException>(op.BytesTransferred <= op.Buffer.Length);
			try
			{
				DeserializeMessageAndQueueForReceiveNotification(op.RemoteEndPoint, op.Buffer, 0, op.BytesTransferred);
			}
			finally
			{
				RecycleBuffer(op.Buffer);
				op.EndAsyncOp();
			}
		}

		void DeserializeMessageAndQueueForReceiveNotification(EndPoint rcep
			, byte[] buffer
			, int offset, int size)
		{
			Contract.Requires<ArgumentNullException>(rcep != null);
			Contract.Requires<ArgumentNullException>(buffer != null);
			Contract.Requires<ArgumentOutOfRangeException>(offset >= 0);
			Contract.Requires<ArgumentOutOfRangeException>(size >= 0);
			Contract.Requires<ArgumentOutOfRangeException>(size <= (buffer.Length - offset));

			var ep = (IPEndPoint) rcep;
			var timestamp = DateTime.UtcNow;
			var decode = _protocol.TryDecodeMessage(ep, buffer, offset, size);
			switch (decode.Kind)
			{
				case DecodeMessageResultKind.Success:
					_notifier.Push(new InboundMessage(timestamp, decode.Message, ep));
					break;
				case DecodeMessageResultKind.Partial:
					// TODO: implement partial message (overrun)
					break;
				case DecodeMessageResultKind.Garbage:
					// TODO: Implement garbage traffic handling.
					break;
				case DecodeMessageResultKind.Blacklist:
					// TODO: Implement blacklist handling.
					break;
			}
		}

		/// <summary>
		/// Structure encapsulating an inbound message.
		/// </summary>
		protected struct InboundMessage
		{
			/// <summary>
			/// The endpoint where the message arrived.
			/// </summary>
			public readonly IPEndPoint Endpoint;
			/// <summary>
			/// The deserialized message.
			/// </summary>
			public readonly TMessage Message;
			/// <summary>
			/// A locally generated UTC timestamp taken when the message arrived.
			/// </summary>
			public readonly DateTime Timestamp;

			/// <summary>
			/// Creates an instance.
			/// </summary>
			/// <param name="timestamp"></param>
			/// <param name="message"></param>
			/// <param name="ep"></param>
			public InboundMessage(DateTime timestamp, TMessage message, IPEndPoint ep)
			{
				this.Timestamp = timestamp;
				this.Message = message;
				this.Endpoint = ep;
			}
		}
	}
}